#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
websocket 链接管理
"""

from typing import List, Dict

from starlette.websockets import WebSocket


class ConnectionManager:

    def __init__(self):
        """存放链接"""
        self.active_connections: List[Dict[str, WebSocket]] = []

    async def connect(self, user: str, ws: WebSocket):
        """链接"""
        await ws.accept()
        self.active_connections.append({"user": user, "ws": ws})

    async def disconnect(self, user: str, ws: WebSocket):
        """断开链接，移除"""
        self.active_connections.remove({"user": user, "ws": ws})

    @staticmethod
    async def send_personal_message(message: str, ws: WebSocket):
        """发送所有人消息"""
        await ws.send_text(message)


    async def send_other_message(self, message: dict, user: str):
        """发送个人消息"""
        for conn in self.active_connections:
            if conn["user"] == user:
                await conn["ws"].send_json(message)


    async def broadcast(self, data: str):
        """广播"""
        for conn in self.active_connections:
            await conn["ws"].send_text(data)


    async def broadcast_other(self, message: str, user: str):
        """广播(排除个人)"""
        for conn in self.active_connections:
            if conn["user"] != user:
                await conn["ws"].send_text(message)


    async def broadcast_json(self, data: dict):
        """广播json数据"""
        for conn in self.active_connections:
            await conn["ws"].send_json(data)

    @staticmethod
    async def send_personal_message_json(data: dict, ws: WebSocket):
        """发送所有人消息(json数据格式)"""
        await ws.send_json(data)






